import threading
import time
from django.db import connection
from concurrent.futures import ThreadPoolExecutor
from dvadmin.system.core.wos_download import crawl_doc_info


# 下载任务管理类 单例
class DownloadTaskManager:
    task_map = {}
    lock = threading.Lock()

    """
    num:    下载线程最大容量
    lazy:   是否懒加载（当没有空闲智能体才会初始化新的智能体 否则预先初始化完所有智能体）
    """

    def __init__(self, num: int = 10):
        self.executor = ThreadPoolExecutor(max_workers=num)

    # 部署爬取任务
    def submit_tasks(self, args: list[dict]) -> bool:
        try:
            for arg in args:
                self.executor.submit(crawl_doc_info, arg)
            return True
        except Exception as e:
            print(f'submit_crawl_tasks error: {e}')
            return False


# 智能体
class Agent:
    """
    主要负责为爬取指定年份的期刊文献信息内容做准备
    """
    journal = None
    crawled_years = {}

    def __init__(self, agent_id: int):
        self.agent_id = agent_id
        self.get_journal()

    # 在数据库智能体表中创建新智能体
    def get_journal(self):
        cursor = connection.cursor()
        # 查询数据库中是否存在该期刊
        cursor.execute(f"SELECT journal FROM agent WHERE id={self.agent_id}")
        self.journal = cursor.fetchone()[0]
        #self.journal = models.Agent.objects.filter(id=self.agent_id).values_list('journal', flat=True).first() 和上一行代码功能一致，不过采用了django的orm
        # 更新状态
        cursor.execute(f"UPDATE agent SET status=1 WHERE id={self.agent_id}")
        connection.commit()
        cursor.close()

    # 更新已经爬取过的年份
    def update_years(self):
        cursor = connection.cursor()
        cursor.execute(f"SELECT year,last_query_num FROM agent_record WHERE agent_id={self.agent_id}")
        # 将查询到的 year 和 last_query_num 组成 year:last_query_num 形式的字典
        self.crawled_years = {str(row[0]): row[1] for row in cursor.fetchall()}
        connection.commit()
        cursor.close()

    # 给没有爬取过的年份创建智能体爬取记录
    def get_new_args(self, years: list[str]) -> list[dict]:
        args = []
        cur_time = int(time.time_ns() / 1e9)
        # 定位到对应年份的数据表
        cursor = connection.cursor()
        for year in years:
            cursor.execute(f'CREATE TABLE IF NOT EXISTS document_{year} LIKE document')

            # 写入爬取记录
            cursor.execute(
                f"INSERT INTO agent_record VALUE (DEFAULT,{self.agent_id},{year},{cur_time},0)")

            # 获取记录ID
            cursor.execute("SELECT LAST_INSERT_ID()")
            args.append({
                'journal': self.journal,
                'year': year,
                'agent_id': self.agent_id,
                'agent_record_id': cursor.fetchone()[0],
                'last_num': 0
            })

        # 提交事务 和 关闭游标
        connection.commit()
        cursor.close()
        return args

    def get_old_args(self, years: list[str]) -> list[dict]:
        args = []
        cursor = connection.cursor()
        for year in years:
            cursor.execute(
                f"SELECT id, last_query_num FROM agent_record WHERE agent_id={self.agent_id} AND year={year}")
            result = cursor.fetchone()
            args.append({
                'journal': self.journal,
                'year': year,
                'agent_id': self.agent_id,
                'agent_record_id': result[0],
                'last_num': max(result[1], 0)
            })
        connection.commit()
        cursor.close()
        return args

    def add_tasks(self, manager: DownloadTaskManager, years: list[str]) -> bool:
        # 更新当前年份信息
        self.update_years()
        not_crawled_years = []
        crawled_years = []

        for year in years:
            if year not in self.crawled_years:
                not_crawled_years.append(year)
            else:
                self.crawled_years[year] = True
                crawled_years.append(year)
        # print(f'{self.crawled_years}\n{not_crawled_years}\n{crawled_years}')

        # 提交新爬取任务和更新任务
        return manager.submit_tasks(self.get_new_args(not_crawled_years) + self.get_old_args(crawled_years))

    def __del__(self):
        cursor = connection.cursor()
        # 更新状态
        cursor.execute(f"UPDATE agent SET status=0 WHERE id={self.agent_id}")
        cursor.execute(f"UPDATE agent SET is_running=FALSE WHERE id={self.agent_id}")
        connection.commit()
        cursor.close()

